import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# Pinecone API key 和配置
api_key = "e88de1a8-5670-4f78-8466-eaa46d703a9a"  # 替换为你的Pinecone API密钥
pinecone = Pinecone(api_key=api_key)  # 连接到Pinecone

# 索引名称
index_name = "mnist-index"  # 替换为你的Pinecone索引名称

# 连接到索引
index = pinecone.Index(index_name)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    try:
        # 确保输入是 PIL 图像
        if not isinstance(image, Image.Image):
            return "Error: Input is not a PIL image."

        # 将输入的手写图像转换为灰度图，并调整大小为8x8
        image = image.convert('L').resize((8, 8))

        # 将图像数据转换为数组，并将其展平为1D数组
        img_array = np.array(image).reshape(1, -1)

        # 对图像数据进行归一化处理，使其范围为 0 到 1
        img_array_normalized = img_array / 255.0

        # 使用 Pinecone 进行搜索
        results = index.query(
            vector=img_array_normalized.ravel().tolist(),
            top_k=11,  # 返回距离最近的 11 个结果
            include_metadata=True  # 同时返回每个向量的元数据(包括标签)
        )

        # 从搜索结果中提取每个匹配项的标签
        labels = [match['metadata']['label'] for match in results['matches']]

        # 使用投票机制确定最终的分类结果
        if labels:
            final_prediction = Counter(labels).most_common(1)[0][0]
        else:
            final_prediction = "No match found"

        return final_prediction
    except Exception as e:
        return f"Error during prediction: {e}"

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,  # 调用的函数
    inputs=gr.Sketchpad(type="pil"),  # 手写板输入，类型为 PIL 图像
    outputs=gr.Label(num_top_classes=1),  # 输出是预测的标签
    title="手写数字识别",  # Web 应用标题
    description="在手写板上输入一个数字，模型将返回预测结果。"  # 描述信息
)

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()